{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "------------ Options -------------\n",
      "appearance_generation: False\n",
      "aspect_ratio: 1.0\n",
      "batchSize: 1\n",
      "checkpoints_dir: ./checkpoints\n",
      "cluster_path: features_clustered_010.npy\n",
      "data_type: 32\n",
      "dataroot: ./datasets/zalando_data/\n",
      "densepose_nc: 3\n",
      "display_winsize: 512\n",
      "engine: None\n",
      "export_onnx: None\n",
      "feat_num: 30\n",
      "fineSize: 512\n",
      "fp16: False\n",
      "gpu_ids: [0]\n",
      "how_many: 50\n",
      "input_nc: 20\n",
      "instance_feat: False\n",
      "isTrain: False\n",
      "label_feat: False\n",
      "label_nc: 20\n",
      "loadSize: 1024\n",
      "load_features: False\n",
      "local_rank: 0\n",
      "max_dataset_size: inf\n",
      "model: ov_pix2pixHD\n",
      "nThreads: 2\n",
      "n_blocks_global: 9\n",
      "n_blocks_local: 3\n",
      "n_clusters: 10\n",
      "n_downsample_E: 4\n",
      "n_downsample_global: 4\n",
      "n_local_enhancers: 1\n",
      "name: zalando\n",
      "nef: 16\n",
      "netG: global\n",
      "ngf: 64\n",
      "niter_fix_global: 0\n",
      "no_flip: False\n",
      "no_instance: False\n",
      "norm: instance\n",
      "ntest: inf\n",
      "onnx: None\n",
      "output_nc: 20\n",
      "phase: test\n",
      "resize_or_crop: scale_width\n",
      "results_dir: ./results/\n",
      "serial_batches: False\n",
      "shape_generation: False\n",
      "tf_log: False\n",
      "use_dropout: False\n",
      "use_encoded_image: False\n",
      "use_generator_last_activation: False\n",
      "verbose: False\n",
      "which_epoch: latest\n",
      "-------------- End ----------------\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from collections import OrderedDict\n",
    "from torch.autograd import Variable\n",
    "from options.test_options import TestOptions\n",
    "from data.ov_test_dataset import TestDataset\n",
    "from models.models import create_model\n",
    "import util.util as util\n",
    "from util.visualizer import Visualizer\n",
    "from util import html\n",
    "import torch\n",
    "from torchvision import transforms\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "opt = TestOptions().parse(save=False)\n",
    "opt.nThreads = 1   # test code only supports nThreads = 1\n",
    "opt.batchSize = 1  # test code only supports batchSize = 1\n",
    "opt.serial_batches = True  # no shuffle\n",
    "opt.no_flip = True  # no flip\n",
    "\n",
    "opt.label_feat = True\n",
    "opt.model = 'ov_pix2pixHD'\n",
    "opt.name = 'zalando_shape'\n",
    "opt.dataroot= './datasets/test_data/'\n",
    "opt.shape_generation = True\n",
    "# opt.use_generator_last_activation = False\n",
    "opt.input_nc= 20 \n",
    "opt.output_nc = 20"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#testing images = 12\n"
     ]
    }
   ],
   "source": [
    "# data_loader = CreateDataLoader(opt)\n",
    "# dataset = data_loader.load_data()\n",
    "augment = {}\n",
    "augment['1'] = transforms.Compose(\n",
    "    [\n",
    "        transforms.ToPILImage(),\n",
    "        transforms.RandomAffine(degrees=10, translate=(0.1, 0.1), scale=(0.8, 1.2), shear=20),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.5, ), (0.5, ))])  # change to [C, H, W]\n",
    "\n",
    "augment['2'] = transforms.Compose(\n",
    "    [\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.5, ), (0.5, ))])  # change to [C, H, W]\n",
    "\n",
    "test_dataset = TestDataset(opt, augment)\n",
    "\n",
    "test_dataloader = DataLoader(test_dataset,\n",
    "                              shuffle=False,\n",
    "                              drop_last=False,\n",
    "                              num_workers=6,\n",
    "                              batch_size=opt.batchSize,\n",
    "                              pin_memory=True)\n",
    "dataset_size = len(test_dataset)\n",
    "print('#testing images = %d' % dataset_size)\n",
    "\n",
    "## FOR DEBUGGING \n",
    "# input_dict = {\n",
    "#             'query_parse_map': A_tensor_label,\n",
    "#             'ref_parse_map': B_tensor_label,\n",
    "#             'query_seg_map': query_label_seg_mask,\n",
    "#             'ref_seg_map': ref_label_seg_mask,\n",
    "#             'query_img': A_tensor_img,\n",
    "#             'ref_img': B_tensor_img,\n",
    "#             'C_tensor': C_tensor\n",
    "#         }\n",
    "# for i in test_dataset[0].keys():\n",
    "#     try:\n",
    "#         print('{}-------'.format(i),test_dataset[0][i].shape)\n",
    "#     except Exception as e:\n",
    "#         print(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------- model used --------- ov_pix2pixHD\n",
      "GlobalGenerator(\n",
      "  (model): Sequential(\n",
      "    (0): ReflectionPad2d((3, 3, 3, 3))\n",
      "    (1): Conv2d(227, 64, kernel_size=(7, 7), stride=(1, 1))\n",
      "    (2): InstanceNorm2d(64, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "    (3): ReLU(inplace=True)\n",
      "    (4): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "    (5): InstanceNorm2d(128, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "    (6): ReLU(inplace=True)\n",
      "    (7): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "    (8): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "    (9): ReLU(inplace=True)\n",
      "    (10): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "    (11): InstanceNorm2d(512, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "    (12): ReLU(inplace=True)\n",
      "    (13): Conv2d(512, 1024, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "    (14): InstanceNorm2d(1024, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "    (15): ReLU(inplace=True)\n",
      "    (16): ResnetBlock(\n",
      "      (conv_block): Sequential(\n",
      "        (0): ReflectionPad2d((1, 1, 1, 1))\n",
      "        (1): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1))\n",
      "        (2): InstanceNorm2d(1024, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "        (3): ReLU(inplace=True)\n",
      "        (4): ReflectionPad2d((1, 1, 1, 1))\n",
      "        (5): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1))\n",
      "        (6): InstanceNorm2d(1024, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "      )\n",
      "    )\n",
      "    (17): ResnetBlock(\n",
      "      (conv_block): Sequential(\n",
      "        (0): ReflectionPad2d((1, 1, 1, 1))\n",
      "        (1): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1))\n",
      "        (2): InstanceNorm2d(1024, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "        (3): ReLU(inplace=True)\n",
      "        (4): ReflectionPad2d((1, 1, 1, 1))\n",
      "        (5): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1))\n",
      "        (6): InstanceNorm2d(1024, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "      )\n",
      "    )\n",
      "    (18): ResnetBlock(\n",
      "      (conv_block): Sequential(\n",
      "        (0): ReflectionPad2d((1, 1, 1, 1))\n",
      "        (1): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1))\n",
      "        (2): InstanceNorm2d(1024, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "        (3): ReLU(inplace=True)\n",
      "        (4): ReflectionPad2d((1, 1, 1, 1))\n",
      "        (5): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1))\n",
      "        (6): InstanceNorm2d(1024, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "      )\n",
      "    )\n",
      "    (19): ResnetBlock(\n",
      "      (conv_block): Sequential(\n",
      "        (0): ReflectionPad2d((1, 1, 1, 1))\n",
      "        (1): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1))\n",
      "        (2): InstanceNorm2d(1024, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "        (3): ReLU(inplace=True)\n",
      "        (4): ReflectionPad2d((1, 1, 1, 1))\n",
      "        (5): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1))\n",
      "        (6): InstanceNorm2d(1024, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "      )\n",
      "    )\n",
      "    (20): ResnetBlock(\n",
      "      (conv_block): Sequential(\n",
      "        (0): ReflectionPad2d((1, 1, 1, 1))\n",
      "        (1): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1))\n",
      "        (2): InstanceNorm2d(1024, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "        (3): ReLU(inplace=True)\n",
      "        (4): ReflectionPad2d((1, 1, 1, 1))\n",
      "        (5): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1))\n",
      "        (6): InstanceNorm2d(1024, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "      )\n",
      "    )\n",
      "    (21): ResnetBlock(\n",
      "      (conv_block): Sequential(\n",
      "        (0): ReflectionPad2d((1, 1, 1, 1))\n",
      "        (1): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1))\n",
      "        (2): InstanceNorm2d(1024, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "        (3): ReLU(inplace=True)\n",
      "        (4): ReflectionPad2d((1, 1, 1, 1))\n",
      "        (5): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1))\n",
      "        (6): InstanceNorm2d(1024, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "      )\n",
      "    )\n",
      "    (22): ResnetBlock(\n",
      "      (conv_block): Sequential(\n",
      "        (0): ReflectionPad2d((1, 1, 1, 1))\n",
      "        (1): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1))\n",
      "        (2): InstanceNorm2d(1024, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "        (3): ReLU(inplace=True)\n",
      "        (4): ReflectionPad2d((1, 1, 1, 1))\n",
      "        (5): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1))\n",
      "        (6): InstanceNorm2d(1024, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "      )\n",
      "    )\n",
      "    (23): ResnetBlock(\n",
      "      (conv_block): Sequential(\n",
      "        (0): ReflectionPad2d((1, 1, 1, 1))\n",
      "        (1): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1))\n",
      "        (2): InstanceNorm2d(1024, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "        (3): ReLU(inplace=True)\n",
      "        (4): ReflectionPad2d((1, 1, 1, 1))\n",
      "        (5): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1))\n",
      "        (6): InstanceNorm2d(1024, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "      )\n",
      "    )\n",
      "    (24): ResnetBlock(\n",
      "      (conv_block): Sequential(\n",
      "        (0): ReflectionPad2d((1, 1, 1, 1))\n",
      "        (1): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1))\n",
      "        (2): InstanceNorm2d(1024, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "        (3): ReLU(inplace=True)\n",
      "        (4): ReflectionPad2d((1, 1, 1, 1))\n",
      "        (5): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1))\n",
      "        (6): InstanceNorm2d(1024, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "      )\n",
      "    )\n",
      "    (25): ConvTranspose2d(1024, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))\n",
      "    (26): InstanceNorm2d(512, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "    (27): ReLU(inplace=True)\n",
      "    (28): ConvTranspose2d(512, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))\n",
      "    (29): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "    (30): ReLU(inplace=True)\n",
      "    (31): ConvTranspose2d(256, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))\n",
      "    (32): InstanceNorm2d(128, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "    (33): ReLU(inplace=True)\n",
      "    (34): ConvTranspose2d(128, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))\n",
      "    (35): InstanceNorm2d(64, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "    (36): ReLU(inplace=True)\n",
      "    (37): ReflectionPad2d((3, 3, 3, 3))\n",
      "    (38): Conv2d(64, 20, kernel_size=(7, 7), stride=(1, 1))\n",
      "  )\n",
      ")\n",
      "Encoder(\n",
      "  (model): Sequential(\n",
      "    (0): ReflectionPad2d((3, 3, 3, 3))\n",
      "    (1): Conv2d(1, 16, kernel_size=(7, 7), stride=(1, 1))\n",
      "    (2): InstanceNorm2d(16, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "    (3): ReLU(inplace=True)\n",
      "    (4): Conv2d(16, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "    (5): InstanceNorm2d(32, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "    (6): ReLU(inplace=True)\n",
      "    (7): Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "    (8): InstanceNorm2d(64, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "    (9): ReLU(inplace=True)\n",
      "    (10): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "    (11): InstanceNorm2d(128, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "    (12): ReLU(inplace=True)\n",
      "    (13): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
      "    (14): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "    (15): ReLU(inplace=True)\n",
      "    (16): ConvTranspose2d(256, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))\n",
      "    (17): InstanceNorm2d(128, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "    (18): ReLU(inplace=True)\n",
      "    (19): ConvTranspose2d(128, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))\n",
      "    (20): InstanceNorm2d(64, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "    (21): ReLU(inplace=True)\n",
      "    (22): ConvTranspose2d(64, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))\n",
      "    (23): InstanceNorm2d(32, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "    (24): ReLU(inplace=True)\n",
      "    (25): ConvTranspose2d(32, 16, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))\n",
      "    (26): InstanceNorm2d(16, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
      "    (27): ReLU(inplace=True)\n",
      "    (28): ReflectionPad2d((3, 3, 3, 3))\n",
      "    (29): Conv2d(16, 10, kernel_size=(7, 7), stride=(1, 1))\n",
      "    (30): Tanh()\n",
      "  )\n",
      ")\n",
      "netG_input_nc, netG_output_nc 227 20\n",
      "netE_input_nc, netE_output_nc  1 10\n",
      "./checkpoints/zalando_shape/latest_net_G.pth\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "./checkpoints/zalando_shape/latest_net_E.pth\n"
     ]
    }
   ],
   "source": [
    "model = create_model(opt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "# torch.unique(test_dataset[0]['C_tensor'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from matplotlib import pyplot as plt\n",
    "# from PIL import Image\n",
    "# from PIL import Image\n",
    "# import numpy as np\n",
    "# x = test_dataset[0]['C_tensor'].cpu().float().numpy()\n",
    "# for i in range(20):\n",
    "#     plt.show()\n",
    "#     plt.imshow(Image.fromarray((x[i,:,:]).astype(np.uint8)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "process image... ['./datasets/test_data/test_query_label/00_Y0121E0K6-J11@9.png']\n",
      "process image... ['./datasets/test_data/test_query_label/0_4BE21E07Q-K11@10.png']\n",
      "process image... ['./datasets/test_data/test_query_label/0_N00098.png']\n",
      "process image... ['./datasets/test_data/test_query_label/1_B0N21E061-Q11@7.png']\n",
      "process image... ['./datasets/test_data/test_query_label/1_N00155.png']\n",
      "process image... ['./datasets/test_data/test_query_label/2_AM421E00G-K11@8.png']\n",
      "process image... ['./datasets/test_data/test_query_label/2_N00383.png']\n",
      "process image... ['./datasets/test_data/test_query_label/3_31021E00B-A11@8.png']\n",
      "process image... ['./datasets/test_data/test_query_label/3_N00807.png']\n",
      "process image... ['./datasets/test_data/test_query_label/4_AM421E00N-A11@11.png']\n",
      "process image... ['./datasets/test_data/test_query_label/4_N01497.png']\n",
      "process image... ['./datasets/test_data/test_query_label/5_ARC21E00O-E11@14.png']\n"
     ]
    }
   ],
   "source": [
    "visualizer = Visualizer(opt)\n",
    "# create website\n",
    "web_dir = os.path.join(opt.results_dir, opt.name, '%s_%s' % (opt.phase, opt.which_epoch))\n",
    "webpage = html.HTML(web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' % (opt.name, opt.phase, opt.which_epoch))\n",
    "for i, data in enumerate(test_dataloader):\n",
    "    if i >= opt.how_many:\n",
    "        break\n",
    "    query_ref_mixed, generated = model.inference_forward_shape(data['query_parse_map'],data['ref_parse_map'],\n",
    "                                                              data['C_tensor'])\n",
    "    visuals = OrderedDict([('query', util.tensor2label(data['query_parse_map'][0], opt.label_nc)),\n",
    "                           ('ref', util.tensor2label(data['ref_parse_map'][0], opt.label_nc)),\n",
    "                           ('query_ref_mixed', util.tensor2label(query_ref_mixed.data[0], opt.label_nc)),\n",
    "                           ('synthesized_Simage', util.tensor2label(generated.data[0], opt.label_nc)),\n",
    "                          ('synthesized_image_edgemap', util.tensor2edgemap(torch.softmax(generated.data[0],dim=0)))])\n",
    "    img_path = data['path']\n",
    "    print('process image... %s' % img_path)\n",
    "    visualizer.save_images(webpage, visuals, img_path)\n",
    "\n",
    "webpage.save()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import numpy as np\n",
    "# np.unique(visuals['synthesized_image_edgemap'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from PIL import Image\n",
    "# test = Image.open('./datasets/fashion_compatibility/test_query_ref_label/0_N00098_synthesized_image_edgemap.jpg')\n",
    "# print(np.unique(np.array(test)))\n",
    "# test = Image.open('./results/fashion_compatibility_shape/test_latest/images/0_N00098_synthesized_image_edgemap.png')\n",
    "# np.unique(np.array(test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from collections import OrderedDict\n",
    "from torch.autograd import Variable\n",
    "from options.test_options import TestOptions\n",
    "from data.ov_test_dataset import TestDataset\n",
    "from models.models import create_model\n",
    "import util.util as util\n",
    "from util.visualizer import Visualizer\n",
    "from util import html\n",
    "import torch\n",
    "from torchvision import transforms\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "opt = TestOptions().parse(save=False)\n",
    "opt.nThreads = 1   # test code only supports nThreads = 1\n",
    "opt.batchSize = 1  # test code only supports batchSize = 1\n",
    "opt.serial_batches = True  # no shuffle\n",
    "opt.no_flip = True  # no flip\n",
    "\n",
    "opt.label_feat = True\n",
    "opt.model = 'ov_pix2pixHD'\n",
    "opt.name = 'zalando_appearance'\n",
    "# opt.dataroot= './datasets/fashion_compatibility/'\n",
    "opt.appearance_generation = True\n",
    "opt.use_generator_last_activation = True\n",
    "opt.input_nc= 20 \n",
    "opt.output_nc = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# data_loader = CreateDataLoader(opt)\n",
    "# dataset = data_loader.load_data()\n",
    "augment = {}\n",
    "augment['1'] = transforms.Compose(\n",
    "    [\n",
    "        transforms.ToPILImage(),\n",
    "        transforms.RandomAffine(degrees=10, translate=(0.1, 0.1), scale=(0.8, 1.2), shear=20),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.5, ), (0.5, ))])  # change to [C, H, W]\n",
    "\n",
    "augment['2'] = transforms.Compose(\n",
    "    [\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.5, ), (0.5, ))])  # change to [C, H, W]\n",
    "\n",
    "test_dataset = TestDataset(opt, augment)\n",
    "\n",
    "test_dataloader = DataLoader(test_dataset,\n",
    "                              shuffle=False,\n",
    "                              drop_last=False,\n",
    "                              num_workers=6,\n",
    "                              batch_size=opt.batchSize,\n",
    "                              pin_memory=True)\n",
    "dataset_size = len(test_dataset)\n",
    "print('#testing images = %d' % dataset_size)\n",
    "\n",
    "# # FOR DEBUGGING \n",
    "# input_dict = {\n",
    "#             'query_parse_map': A_tensor_label,\n",
    "#             'ref_parse_map': B_tensor_label,\n",
    "#             'query_seg_map': query_label_seg_mask,\n",
    "#             'ref_seg_map': ref_label_seg_mask,\n",
    "#             'query_img': A_tensor_img,\n",
    "#             'ref_img': B_tensor_img,\n",
    "#             'C_tensor': C_tensor\n",
    "#         }\n",
    "# for i in test_dataset[0].keys():\n",
    "#     try:\n",
    "#         print('{}-------'.format(i),test_dataset[0][i].shape)\n",
    "#     except Exception as e:\n",
    "#         print(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = create_model(opt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualizer = Visualizer(opt)\n",
    "# create website\n",
    "web_dir = os.path.join(opt.results_dir, opt.name, '%s_%s' % (opt.phase, opt.which_epoch))\n",
    "webpage = html.HTML(web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' % (opt.name, opt.phase, opt.which_epoch))\n",
    "for i, data in enumerate(test_dataloader):\n",
    "    if i >= opt.how_many:\n",
    "        break\n",
    "    generated = model.inference_forward_appearance(data['query_img'],data['query_parse_map'],\n",
    "                                                              data['query_seg_map'],data['ref_img'],\n",
    "                                                              data['ref_parse_map'],data['ref_seg_map'],\n",
    "                                                              data['C_tensor'])\n",
    "    \n",
    "    visuals = OrderedDict([('query_img', util.tensor2im(data['query_img'][0])),\n",
    "                           ('ref_image', util.tensor2im(data['ref_img'][0])),\n",
    "                           ('generated_parse_map', util.tensor2label(data['C_tensor'][0], opt.label_nc)),\n",
    "                           ('synthesized_image', util.tensor2im(generated.data[0]))])\n",
    "    \n",
    "    img_path = data['path']\n",
    "    print('process image... %s' % img_path)\n",
    "    visualizer.save_images(webpage, visuals, img_path)\n",
    "\n",
    "webpage.save()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "app_feature_vectors_query = torch.zeros((20,30)).float().cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "app_feature_vectors_query[0,0] = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from collections import OrderedDict\n",
    "from torch.autograd import Variable\n",
    "from options.test_options import TestOptions\n",
    "from data.ov_test_dataset import TestDataset\n",
    "from models.models import create_model\n",
    "import util.util as util\n",
    "from util.visualizer import Visualizer\n",
    "from util import html\n",
    "import torch\n",
    "from torchvision import transforms\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "opt = TestOptions().parse(save=False)\n",
    "opt.nThreads = 1   # test code only supports nThreads = 1\n",
    "opt.batchSize = 1  # test code only supports batchSize = 1\n",
    "opt.serial_batches = True  # no shuffle\n",
    "opt.no_flip = True  # no flip\n",
    "\n",
    "opt.model = 'ov_pix2pixHD'\n",
    "opt.name = 'fashion_compatibility_appearance'\n",
    "# opt.dataroot= './datasets/fashion_compatibility/'\n",
    "opt.appearance_generation = True\n",
    "opt.use_generator_last_activation = True\n",
    "opt.input_nc= 20 \n",
    "opt.output_nc = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# data_loader = CreateDataLoader(opt)\n",
    "# dataset = data_loader.load_data()\n",
    "augment = {}\n",
    "augment['1'] = transforms.Compose(\n",
    "    [\n",
    "        transforms.ToPILImage(),\n",
    "        transforms.RandomAffine(degrees=10, translate=(0.1, 0.1), scale=(0.8, 1.2), shear=20),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.5, ), (0.5, ))])  # change to [C, H, W]\n",
    "\n",
    "augment['2'] = transforms.Compose(\n",
    "    [\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.5, ), (0.5, ))])  # change to [C, H, W]\n",
    "\n",
    "test_dataset = TestDataset(opt, augment)\n",
    "\n",
    "test_dataloader = DataLoader(test_dataset,\n",
    "                              shuffle=False,\n",
    "                              drop_last=False,\n",
    "                              num_workers=6,\n",
    "                              batch_size=opt.batchSize,\n",
    "                              pin_memory=True)\n",
    "dataset_size = len(test_dataset)\n",
    "print('#testing images = %d' % dataset_size)\n",
    "\n",
    "# # FOR DEBUGGING \n",
    "# input_dict = {\n",
    "#             'query_parse_map': A_tensor_label,\n",
    "#             'ref_parse_map': B_tensor_label,\n",
    "#             'query_seg_map': query_label_seg_mask,\n",
    "#             'ref_seg_map': ref_label_seg_mask,\n",
    "#             'query_img': A_tensor_img,\n",
    "#             'ref_img': B_tensor_img,\n",
    "#             'C_tensor': C_tensor\n",
    "#         }\n",
    "# for i in test_dataset[0].keys():\n",
    "#     try:\n",
    "#         print('{}-------'.format(i),test_dataset[0][i].shape)\n",
    "#     except Exception as e:\n",
    "#         print(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from collections import OrderedDict\n",
    "from torch.autograd import Variable\n",
    "from options.train_options import TrainOptions\n",
    "from data.ov_test_dataset import TestDataset\n",
    "from models.models import create_model\n",
    "import util.util as util\n",
    "from util.visualizer import Visualizer\n",
    "from util import html\n",
    "import torch\n",
    "from torchvision import transforms\n",
    "from torch.utils.data import DataLoader\n",
    "import time\n",
    "\n",
    "opt = TrainOptions().parse(save=False)\n",
    "opt.nThreads = 1   # test code only supports nThreads = 1\n",
    "opt.batchSize = 1  # test code only supports batchSize = 1\n",
    "opt.serial_batches = True  # no shuffle\n",
    "opt.no_flip = True  # no flip\n",
    "\n",
    "opt.model = 'ov_pix2pixHD_online'\n",
    "opt.name = 'fashion_compatibility_online'\n",
    "# opt.dataroot= './datasets/fashion_compatibility/'\n",
    "opt.appearance_generation = True\n",
    "opt.use_generator_last_activation = True\n",
    "opt.input_nc= 20 \n",
    "opt.output_nc = 3\n",
    "opt.load_pretrain = './checkpoints/zalando_appearance/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = create_model(opt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training Visualizer\n",
    "visualizer = Visualizer(opt)\n",
    "\n",
    "# Optimizers\n",
    "optimizer_G, optimizer_D = model.module.optimizer_G, model.module.optimizer_D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prev_j = 0\n",
    "index_check = 20\n",
    "for i, data in enumerate(test_dataloader, start=0):\n",
    "    for j in range(0,100):\n",
    "        save_fake = True if j % index_check == 0 else False\n",
    "        losses, generated = model(data['query_img'],data['query_parse_map'],\n",
    "                                                    data['query_seg_map'],data['ref_img'],\n",
    "                                                    data['ref_parse_map'],data['ref_seg_map'],\n",
    "                                                    data['C_tensor'],infer=save_fake)\n",
    "\n",
    "        # sum per device losses\n",
    "        losses = [torch.mean(x) if not isinstance(x, int) else x for x in losses]\n",
    "        loss_dict = dict(zip(model.module.loss_names, losses))\n",
    "\n",
    "        # calculate final loss scalar\n",
    "        loss_ref = (loss_dict['G_GAN_REF'] + loss_dict['G_VGG'])\n",
    "        loss_query = loss_dict['G_GAN_QUERY']\n",
    "        loss_online = loss_ref + loss_query\n",
    "\n",
    "        ############### Backward Pass ####################\n",
    "        # update generator weights\n",
    "        optimizer_G.zero_grad()\n",
    "        loss_online.backward()\n",
    "        optimizer_G.step()\n",
    "        \n",
    "        print('loss_ref {}, loss_query {}, loss_online {}'.format(loss_ref,loss_query,loss_online))\n",
    "        \n",
    "        ############## Display results and errors ##########\n",
    "        if save_fake:\n",
    "            visuals = OrderedDict([('query_img', util.tensor2im(data['query_img'][0])),\n",
    "                               ('ref_image', util.tensor2im(data['ref_img'][0])),\n",
    "                               ('synthesized_image', util.tensor2im(generated.data[0]))])\n",
    "            visualizer.display_current_results(visuals, (j // index_check) + prev_j, i)\n",
    "    prev_j = (j // index_check)*(i+1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test exercises"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_numpy = test_dataset[2]['ref_img'].cpu().float().numpy()\n",
    "image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0\n",
    "image_numpy = image_numpy.astype(np.uint8)\n",
    "Image.fromarray(image_numpy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "org_seg_map = test_dataset[2]['ref_parse_map'].clone()\n",
    "\n",
    "# unnormalize it \n",
    "unnormalize_seg_map = org_seg_map.cpu().float().numpy()\n",
    "unnormalize_seg_map = (np.transpose(unnormalize_seg_map, (1, 2, 0)) + 1) / 2.0\n",
    "unnormalize_seg_map = unnormalize_seg_map.astype(np.uint8)\n",
    "np.unique(unnormalize_seg_map)\n",
    "\n",
    "# filter the top section\n",
    "filter_part = unnormalize_seg_map[:, :, 5] + unnormalize_seg_map[:, :, 6] + unnormalize_seg_map[ :, :,7]  # 5,6,7,\n",
    "filter_part = filter_part > 0\n",
    "filter_part = np.expand_dims(filter_part,-1)\n",
    "\n",
    "\n",
    "image_numpy = test_dataset[2]['ref_img'].cpu().float().numpy()\n",
    "image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0\n",
    "image_numpy = image_numpy.astype(np.uint8)\n",
    "source_img = image_numpy * filter_part # other part normalized 0 --> 127 gray\n",
    "Image.fromarray(source_img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ref_parse_map_online = torch.zeros((20, 512, 256)).float().cuda()\n",
    "for num_seg_channel in range(20):\n",
    "    if 4 < num_seg_channel < 8:\n",
    "        ref_parse_map_online[num_seg_channel,:,:] = test_dataset[2]['ref_parse_map'][num_seg_channel,:,:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ref_parse_map_online = (ref_parse_map_online + 1)/2\n",
    "filter_part = ref_parse_map_online[5, :, :] + ref_parse_map_online[6, :, :] + ref_parse_map_online[7, :, :]  # 5,6,7,\n",
    "filter_part = filter_part > 0\n",
    "filter_part = torch.unsqueeze(filter_part,0).float().cuda()\n",
    "src = test_dataset[2]['ref_img'].float().cuda()\n",
    "src = (src + 1)/2\n",
    "source_img = src * filter_part # other part normalized 0 --> 127 gray\n",
    "test = source_img.clone()\n",
    "image_numpy = source_img.cpu().float().numpy()\n",
    "image_numpy = np.transpose(image_numpy,(1, 2, 0))* 255.0\n",
    "image_numpy = image_numpy.astype(np.uint8)\n",
    "Image.fromarray(image_numpy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "norm = transforms.Normalize((0.5,), (0.5,))\n",
    "out = norm(test)\n",
    "torch.unique(out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "pytorch"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
